import os
import onnx
import copy
import numpy as np
import logging
import onnxruntime
import sys, getopt
import json
from onnx import helper

from collections import OrderedDict
from onnx import shape_inference

logging.basicConfig(level=logging.INFO)

from onnx import shape_inference, TensorProto, version_converter, numpy_helper

logger = logging.getLogger("[ONNXOPTIMIZER]")

check_op_list = ['Max', 'Min', 'Sum', 'Mean']

onnxfile = ''
nodefile = ''
outputfile = ''

def model_post_process(onnxfile):

    model = onnx.load(onnxfile)

    logger.info("model_post_process")

    delete_node_id = 0
    delete = False
    export_onnx = ''

    for node_id, node in enumerate(model.graph.node):
        #dict={"name":node.name, "input":node.input, "output":node.output}
        print(node_id, ", name:", node.name, ", input:", node.input, ", output:", node.output,  \
                ", op:", node.op_type, ', len(input):', len(node.input))

        if node.op_type in check_op_list and len(node.input) == 1:
            delete_node_id = node_id
            delete = True
            break

    if delete == True:     
        print('delete: ', delete_node_id)
        delete_node = model.graph.node[delete_node_id]
        print('delete node op:', delete_node.op_type)
        next_node = model.graph.node[delete_node_id+1]

        for i, n_ in enumerate(next_node.input):
            #print('next input:', n_)
            if n_ == delete_node.output[0]:
                print('got it:', n_)
                next_node.input[i] = delete_node.input[0]

        model.graph.node.remove(delete_node)
        export_onnx = "./tmp.onnx"
        onnx.checker.check_model(model)
        onnx.save(model, export_onnx)

    return delete, export_onnx    
 
def model_post_process_tmp(onnxfile):

    model = onnx.load(onnxfile)

    logger.info("model_post_process_tmp")

    delete_node_id = 0
    delete = False
    export_onnx = ''

    for node_id, node in enumerate(model.graph.node):
        #dict={"name":node.name, "input":node.input, "output":node.output}
        print(node_id, ", name:", node.name, ", input:", node.input, ", output:", node.output,  \
                ", op:", node.op_type, ', len(input):', len(node.input))

        if node_id == 16:
            print('++++++++++++++++++++++++++', node.name)
            node.output[0] = 'fc2_1/exc_xxx'
        elif node_id == 17:
            print('---------------------', node.name)
            node.input[0] = 'fc2_1/exc_xxx' 
        elif node_id == 32:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc2_2/exc_xxx'
        elif node_id == 33:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc2_2/exc_xxx'
        elif node_id == 48:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc2_3/exc_xxx'
        elif node_id == 49:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc2_3/exc_xxx'
        elif node_id == 65:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc3_1/exc_xxx'
        elif node_id == 66:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc3_1/exc_xxx'
        elif node_id == 81:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc3_2/exc_xxx'
        elif node_id == 82:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc3_2/exc_xxx'  
        elif node_id == 97:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc3_3/exc_xxx'
        elif node_id == 98:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc3_3/exc_xxx'
        elif node_id == 113:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc3_4/exc_xxx'
        elif node_id == 114:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc3_4/exc_xxx' 
        elif node_id == 130:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc4_1/exc_xxx'
        elif node_id == 131:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc4_1/exc_xxx' 
        elif node_id == 146:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc4_2/exc_xxx'
        elif node_id == 147:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc4_2/exc_xxx'
        elif node_id == 162:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc4_3/exc_xxx'
        elif node_id == 163:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc4_3/exc_xxx'    
        elif node_id == 178:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc4_4/exc_xxx'
        elif node_id == 179:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc4_4/exc_xxx'
        elif node_id == 194:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc4_5/exc_xxx'
        elif node_id == 195:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc4_5/exc_xxx'    
        elif node_id == 210:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc4_6/exc_xxx'
        elif node_id == 211:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc4_6/exc_xxx' 
        elif node_id == 227:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc5_1/exc_xxx'
        elif node_id == 228:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc5_1/exc_xxx'   
        elif node_id == 243:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc5_2/exc_xxx'
        elif node_id == 244:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc5_2/exc_xxx' 
        elif node_id == 259:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.output[0] = 'fc5_3/exc_xxx'
        elif node_id == 260:
            print('xxxxxxxxxxxxxxxxxxx', node.name)
            node.input[0] = 'fc5_3/exc_xxx'                                                       

    export_onnx = "./hello.onnx"
    #onnx.checker.check_model(model)
    onnx.save(model, export_onnx)

    return delete, export_onnx    

def usage():
    print('python onnx_debug.py -i <onnxfile>')
    print('or') 
    print('python onnx_debug.py -i <onnxfile> -n <nodefile> -o <outputfile>', '({})'.format('Not yet support for now'))    

def shape_inference_onnx(onnxfile):
    model = onnx.load(onnxfile)
    path = './shape.onnx'  
    onnx.save(onnx.shape_inference.infer_shapes(model), path)

# onnx 插入新的Node
def insert_node(model, insert_node, follow_up_node):
    # 根据插入Node的输出修改后续node的输入
    follow_up_node.input[0] = insert_node.output[0]
    # 找到后续Node的索引位置，并将插入节点插入到graph中
    for follow_up_node_index, _follow_up_node in enumerate(model.graph.node):
        if _follow_up_node == follow_up_node:
            print("follow_up_node_index: ", follow_up_node_index)
            model.graph.node.insert(follow_up_node_index, insert_node)
            break

def gen_new_onnx(onnxfile):
    model = onnx.load(onnxfile)
    path = './shape_new.onnx'

    graph = model.graph
    node  = graph.node
 
    insert_1 = helper.make_node(
        'GlobalAveragePool', # node name
        ['113'],
        ['500']  # outputs
        )

    follow_up_node = None
    for i in range(len(node)):
        if node[i].op_type == 'Conv' and node[i].name == "Conv_61":
            follow_up_node = node[i]
            print('insert a new node----------------')
            break

    insert_node(model, insert_1, follow_up_node)   

    shape = [4]
    data_array = np.array([1,512,13,13], dtype = np.int64)
    new_shape = helper.make_tensor('new_shape', TensorProto.INT64, shape, data_array)
    graph.initializer.append(new_shape)  ###### important!!!!!!!
    insert_2 = helper.make_node(
        'Expand', # node name
        ['500', 'new_shape'],
        ['600']  # outputs
        )

    follow_up_node = None
    for i in range(len(node)):
        if node[i].op_type == 'Conv' and node[i].name == "Conv_61":
            follow_up_node = node[i]
            print('insert a new node+++++++++++++++++++')
            break

    insert_node(model, insert_2, follow_up_node)   

    onnx.checker.check_model(model)
    #onnx.save(model, path)
    onnx.save(onnx.shape_inference.infer_shapes(model), path)

def get_outpus(onnxfile):

    model = onnx.load(onnxfile)

    node_list = []

    logger.info("get_outpus")

    has_global_average_pool = False
    need_convert = False

    for node_id, node in enumerate(model.graph.node):
        print(node_id, ", name:", node.name, ", input:", node.input, ", output:", node.output,  \
                ", op:", node.op_type)

        if node.op_type == 'GlobalAveragePool':
            has_global_average_pool = True 
            dict={'id':node_id, 'input':node.input, 'output':node.output, 'op':node.op_type}   
            node_list.append(dict)

        #node_list.append(dict)

    if has_global_average_pool == True:
        node_list2 = []

        for v in model.graph.value_info:
            input_shape = v.type.tensor_type.shape.dim
            input_shape = [x.dim_value for x in input_shape]

            dict2 = {'name':v.name, 'shape':input_shape}
            node_list2.append(dict2)

            print("+++++++++++ name:", v.name, input_shape)

        for d in node_list:
            if d['op'] == 'GlobalAveragePool':
                print('op id:', d['id'], ', op input:', d['input'])
                for v in node_list2:
                    #print('v.name:', v['name'])
                    if d['input'][0] == v['name']:
                        print('got GlobalAveragePool, shape:', v['shape'])
                        if v['shape'][2] <= 15 and v['shape'][3] <= 15:
                            need_convert = True
                            print('need to convert') 
                            old_node = model.graph.node[d['id']] 
                            model.graph.node.remove(old_node)
                            new_node = onnx.helper.make_node(
                                                    name = '',
                                                    op_type="AveragePool",
                                                    inputs=d['input'],
                                                    outputs=d['output'],
                                                    kernel_shape=[v['shape'][2], v['shape'][3]], 
                                                    )

                            model.graph.node.insert(d['id'], new_node)

            if need_convert == True:
                onnx.checker.check_model(model)
                onnx.save(model, 'ap2.onnx')

def main(argv):
   global onnxfile
   global nodefile
   global outputfile

   try:
      opts, args = getopt.getopt(argv,"hi:n:o:",["onnx=", "nfile=", "ofile="])
   except getopt.GetoptError:
      usage()
      sys.exit(2)
   for opt, arg in opts:
      if opt == '-h':
         usage()
         sys.exit()
      elif opt in ("-i", "--onnx"):
         onnxfile = arg
      elif opt in ("-n", "--nfile"):
         nodefile = arg   
      elif opt in ("-o", "--ofile"):
         outputfile = arg

   if onnxfile == '' :
      usage()
      sys.exit()
         
   print('model file: ', onnxfile)
   print('node file: ', nodefile)
   print('output file: ', outputfile)

   #shape_inference_onnx(onnxfile)

   get_outpus(onnxfile)

   #gen_new_onnx(onnxfile)

   #delete, post_process_file = model_post_process(onnxfile)

   #while delete == True:
   #    delete, post_process_file = model_post_process(post_process_file)

if __name__ == "__main__":
   main(sys.argv[1:])
